import cv2
import numpy as np
import onnxruntime as ort

# 加载 ONNX 模型
ort_session = ort.InferenceSession("runs/detect/train4/weights/best.onnx")

# 图像预处理
def preprocess(img_path):
    img = cv2.imread(img_path)
    h, w = img.shape[:2]
    img = cv2.resize(img, (640, 640))  # 与训练时一致
    img = img / 255.0
    img = img.transpose(2, 0, 1)  # HWC -> CHW
    img = np.expand_dims(img, axis=0).astype(np.float32)
    return img, (h, w)

# 推理函数
def infer(img):
    inputs = {ort_session.get_inputs()[0].name: img}
    outputs = ort_session.run(None, inputs)
    return outputs

# 后处理（简化版）
def postprocess(outputs, orig_shape):
    # 实现框还原、NMS、类别映射等逻辑
    pass

# 主流程
img, orig_shape = preprocess("datasets/coco8/images/train/000000000025.jpg")
outputs = infer(img)
results = postprocess(outputs, orig_shape)
print(results)

